﻿using System;
using System.Collections.Generic;
using Algoritmia.DataStructures.Queues;
using Algoritmia.DataStructures.Trees;
using Algoritmia.Utils;

namespace Algoritmia.Problems.Traversals
{
    public class LevelOrderTreeTraverser<TNode,TResult> : ITreeTraverser<TNode,TResult>
    {
        private readonly Func<IFifo<IRootedTree<TNode>>> _createFifo;

        public LevelOrderTreeTraverser(Func<IFifo<IRootedTree<TNode>>> createFifo = null)
        {
            _createFifo = createFifo ?? (() => AlgoritmiaFactory.Fifos<IRootedTree<TNode>>.Default);
        }


        #region ITreeTraverser<TNode,TResult> Members

        public IEnumerable<TResult> Traverse(IRootedTree<TNode> tree, 
                                             Func<IRootedTree<TNode>, TResult> visitor = null)
        {
            visitor = visitor ?? (subtree => (TResult)(object)subtree.Root);
            var fifo = _createFifo();
            fifo.Push(tree);
            yield return visitor(tree);
            while( fifo.Count > 0)
            {
                var current = fifo.Pop();
                foreach(var subtree in current.SubTrees())
                {
                    fifo.Push(subtree);
                    yield return visitor(subtree);
                }
            }
        }

        #endregion
    }
}
